#include "stdafx.h"
#include "PhotonMap.h"
#include "PMKdTree.h"
#include "knn.h"
#include "PMHash.h"
#include "CudaUtils.h"

namespace PhotonMapping
{

const int PhotonMap::m_kNNBFQueryWidth = 61440;

PhotonMap::PhotonMap(float p_MaxQueryRadius, uintx p_MaxQueryPhotons, uintx p_Technique) 
	: m_Technique(p_Technique)
{
	using namespace Math;

	// intiialize the direction conversion tables
	for (int i=0; i<256; i++)
	{
		float angle = static_cast<float>(i) * (1.0f / 256.0f) * PI;
		m_CosTheta[i] = Cos(angle);
		m_SinTheta[i] = Sin(angle);
		m_CosPhi[i]   = Cos(2.0f * angle);
		m_SinPhi[i]   = Sin(2.0f * angle);
	}
	m_PhotonMap.push_back(Photon());

	m_MaxPhotons = p_MaxQueryPhotons;
	m_QueryRadius = p_MaxQueryRadius;

	// create cuda timers
	CUT_SAFE_CALL(cutCreateTimer(&m_EstimateTimer));
	CUT_SAFE_CALL(cutCreateTimer(&m_PhotonGatherTimer));
	CUT_SAFE_CALL(cutCreateTimer(&m_ConstructionTimer));
}

PhotonMap::~PhotonMap()
{
	CUT_SAFE_CALL(cutDeleteTimer(m_ConstructionTimer));
	CUT_SAFE_CALL(cutDeleteTimer(m_EstimateTimer));
	CUT_SAFE_CALL(cutDeleteTimer(m_PhotonGatherTimer));
}

void PhotonMap::Create(ScenePtr p_Scene, uintx p_NumPhotons, bool p_CausticMap)
{
	using namespace std;
	using namespace Math;

	m_BBMin = Vector3::INFINITY;
	m_BBMax = -Vector3::INFINITY;
	m_PhotonMap.reserve(p_NumPhotons + 100);
	m_NumPhotons = 0;


	////////////////////////////////////////////////////////
	////////////////////////////////////////////////////////
	// TODO: random seed?
	srand(10000);
	////////////////////////////////////////////////////////
	////////////////////////////////////////////////////////

	// build light histogram
	float totalFlux = 0.0f;
	std::vector<float> lightFluxHistogram;
	for(vector<LightPtr>::const_iterator curLight = p_Scene->GetLights().begin(); curLight != p_Scene->GetLights().end(); ++curLight)
	{
		totalFlux += (*curLight)->GetMaxFlux().Average();
		lightFluxHistogram.push_back(totalFlux);

		// if we want a caustic photon map create projection maps here
		if(p_CausticMap)
			(*curLight)->CreateProjectionMap(p_Scene);
	}

	uintx emittedPhotons = 0;
	while(m_PhotonMap.size() <= p_NumPhotons)
	{
		// determine which light source to pick for shooting a ray
		intx curLight = 0;
		float r = Random(1.0f) * totalFlux;
		for(size_t i = 0; i < lightFluxHistogram.size(); ++i)
		{
			if(r <= lightFluxHistogram[i])
			{
				curLight = i;
				break;
			}
		}

		// shoot ray from light
		LightPtr light = p_Scene->GetLights()[curLight];
		Ray photonRay = light->SampleLightRay(Random(1.0f), Random(1.0f));

		// trace photon
		++emittedPhotons;
		const Color photonPower = Color::ONE;
		PhotonTrace(p_Scene, photonRay, photonPower, 0, 0, p_CausticMap);
	}

	// scale according to emitted photons
	const float scale = totalFlux / static_cast<float>(emittedPhotons);
	for(uintx i = 1; i < m_PhotonMap.size(); ++i)
	{
		Color col;
		col.FromRGBE(m_PhotonMap[i].m_Power);
		col *= scale;
		col.ToRGBE(m_PhotonMap[i].m_Power);
	}
}


void PhotonMap::PhotonTrace(ScenePtr p_Scene, const Math::Ray& p_Ray, const Color& p_Power, uintx p_DiffBounces, uintx p_SpecBounces, bool p_CausticMap)
{
	using namespace Math;
	using namespace std;

	if(p_DiffBounces + p_SpecBounces >= MaxTraces)
		return;

	// find closest intersecting object (if any)
	intx hitResult;
	ObjectPtr curObj;
	float dist = INFINITY;
	for(vector<ObjectPtr>::const_iterator i = p_Scene->GetObjects().begin(); i != p_Scene->GetObjects().end(); ++i)
	{
		intx res;
		if((res = (*i)->TestIntersection(p_Ray, dist)))
		{
			curObj = (*i);
			hitResult = res;
		}
	}

	if(curObj)
	{
		// calculate hit point
		Vector3 hitPoint = p_Ray.m_Position + (p_Ray.m_Direction * dist);
		Vector3 hitNormal = curObj->GetNormalAt(hitPoint);

		// get material
		MaterialPtr curMat = p_Scene->GetMaterials()[curObj->GetMaterial()];

		// calc refraction index and adjust normal/dot product values accordingly
		float eta = 1.0f / curMat->GetRefractionIndex();
		float ni = (hitNormal * -p_Ray.m_Direction);
		if(ni < 0.0f)
		{
			eta = curMat->GetRefractionIndex();
			hitNormal = -hitNormal;
			ni = -ni;
		}

		if(p_CausticMap && p_SpecBounces > 0)
			AddPhoton(hitPoint, -p_Ray.m_Direction, p_Power);

		// calc russian roulette scales
		float transmittance = curMat->GetRefraction().Average();
		float specRefl = curMat->GetSpecReflection().Average();
		float diffRefl = curMat->GetDiffReflection().Average();
		float total = transmittance + specRefl + diffRefl;

		// play russian roulette
		float roulette = Random(total);

		// refract?
		if(roulette < diffRefl)
		{
			// terminate if we create a caustic map
			if(p_CausticMap)
				return;
			else if(p_SpecBounces < 1)
				AddPhoton(hitPoint, -p_Ray.m_Direction, p_Power);

			const float phi = PI * 2.0f * Random(1.0f);
			const float theta = ACos(Sqrt(1.0f - Random(1.0f)));

			// make coord frame coefficients (z in normal direction)
			const float x = Cos(phi) * Sin(theta);
			const float y = Sin(phi) * Sin(theta);
			const float z = Cos(theta);

			// make coord frame
			Vector3 normal(hitNormal);
			Vector3 tangent(GetTangent(normal));

			// make vector from frame times coefficients
			Vector3 reflVec = (tangent * x) + (Cross(tangent, normal) * y) + (normal * z);
			reflVec.Normalize();
			Ray reflRay(hitPoint + reflVec * EPSILON, reflVec);

			// calc reflection power
			Color newPower = p_Power * total * curMat->GetDiffReflection() / diffRefl;
			PhotonTrace(p_Scene, reflRay, newPower, p_DiffBounces + 1, p_SpecBounces, p_CausticMap);
		}
		else if(roulette < (diffRefl + specRefl))  // reflect?
		{
			Vector3 reflVec = 2.0f * hitNormal *(ni) + p_Ray.m_Direction;
			reflVec.Normalize();
			Ray reflRay(hitPoint + reflVec * EPSILON, reflVec);
			Color newPower = total * p_Power * curMat->GetSpecReflection() / specRefl;
			PhotonTrace(p_Scene, reflRay, newPower, p_DiffBounces, p_SpecBounces + 1, p_CausticMap);
		}
		else if(roulette < (diffRefl + specRefl + transmittance))
		{
			Vector3 T = eta * p_Ray.m_Direction + (ni * eta - Sqrt(1.0f + eta * eta * (ni * ni - 1.0f))) * hitNormal;
			T.Normalize();
			Color newPower = total * p_Power * curMat->GetRefraction() / transmittance;
			PhotonTrace(p_Scene, Ray(hitPoint + T * EPSILON, T), newPower, p_DiffBounces, p_SpecBounces + 1, p_CausticMap);
		}

		if(!p_CausticMap && p_DiffBounces == 0)
			ShadowPhotonTrace(p_Scene, Ray(p_Ray.m_Position + p_Ray.m_Direction * EPSILON, p_Ray.m_Direction), Color::ZERO);
	}
}

void PhotonMap::ShadowPhotonTrace(ScenePtr p_Scene, const Math::Ray& p_Ray, const Color& p_Power)
{
	using namespace Math;
	using namespace std;

	intx hitResult;
	ObjectPtr curObj;
	float dist = INFINITY;
	for(vector<ObjectPtr>::const_iterator i = p_Scene->GetObjects().begin(); i != p_Scene->GetObjects().end(); ++i)
	{
		intx res;
		if((res = (*i)->TestIntersection(p_Ray, dist)))
		{
			curObj = (*i);
			hitResult = res;
		}
	}

	// if we hit a diffuse object - store a shadow photon there and continue tracing
	if(curObj)
	{
		// get material
		MaterialPtr curMat = p_Scene->GetMaterials()[curObj->GetMaterial()];
		Vector3 hitPosition = p_Ray.m_Position + (p_Ray.m_Direction * dist);
		if(!curMat->GetDiffReflection().IsZero())
			AddPhoton(hitPosition, p_Ray.m_Direction, p_Power);
		
		ShadowPhotonTrace(p_Scene, Ray(hitPosition + p_Ray.m_Direction * EPSILON, p_Ray.m_Direction), p_Power);
	}
}

void PhotonMap::AddPhoton(const Math::Vector3& p_Position, const Math::Vector3& p_Dir, const Color& p_Power)
{
	using namespace Math;

	Photon photon;

	// store position and update bounding box
	photon.m_Position = p_Position;
	m_BBMin.m_X = Min(m_BBMin.m_X, p_Position.m_X);
	m_BBMin.m_Y = Min(m_BBMin.m_Y, p_Position.m_Y);
	m_BBMin.m_Z = Min(m_BBMin.m_Z, p_Position.m_Z);
	m_BBMax.m_X = Max(m_BBMax.m_X, p_Position.m_X);
	m_BBMax.m_Y = Max(m_BBMax.m_Y, p_Position.m_Y);
	m_BBMax.m_Z = Max(m_BBMax.m_Z, p_Position.m_Z);

	// compress power
	//photon.m_Power = p_Power;
	p_Power.ToRGBE(photon.m_Power);

	// calc theta and phi
	int theta = static_cast<int>(ACos(p_Dir.m_Z) * (256.0f / PI));
	if (theta > 255)
		photon.m_Theta = 255;
	else
		photon.m_Theta = static_cast<byte>(theta);

	int phi = static_cast<int>(ATan2(p_Dir.m_Y, p_Dir.m_X) * (256.0f / (2.0f * PI)));
	if (phi > 255)
		photon.m_Phi = 255;
	else if (phi < 0)
		photon.m_Phi = static_cast<byte>(phi + 256);
	else
		photon.m_Phi = static_cast<byte>(phi);

	m_PhotonMap.push_back(photon);
	++m_NumPhotons;
}

void PhotonMap::Prepare()
{
	// select the right technique
	switch(m_Technique)
	{
		case SWKdTree:
			{
				CUT_SAFE_CALL(cutStartTimer(m_ConstructionTimer));
				BuildKdTree();
				CUT_SAFE_CALL(cutStopTimer(m_ConstructionTimer));
			}
			break;

		case GPUKdTree:
			{
				StartAllocationStats();				
				m_KdTree = CreateKdTree(&m_PhotonMap[1], static_cast<unsigned int>(m_PhotonMap.size() - 1), 
					static_cast<unsigned int>(m_MaxPhotons), m_QueryRadius, m_ConstructionTimer);
				EndAllocationStats();
			}
			break;
		case GPUHash:
			{
				StartAllocationStats();
				m_HashMap = PerformHashing(&m_PhotonMap[1], static_cast<unsigned int>(m_PhotonMap.size() - 1), 
					m_BBMin.m_Xyz, m_BBMax.m_Xyz, m_QueryRadius, m_ConstructionTimer);
				EndAllocationStats();
			}
			break;
		case GPUBF:
		default:
			{
				// get photons into right format
				int photonsWidth = static_cast<int>(m_PhotonMap.size() - 1);
				float* photons = new float[photonsWidth * 3];
				float* photonsX = photons;
				float* photonsY = photonsX + photonsWidth;
				float* photonsZ = photonsY + photonsWidth;
				for(int i = 0; i < photonsWidth; ++i)
				{
					const Photon& p = m_PhotonMap[i + 1];
					*photonsX++ = p.m_Position.m_X;
					*photonsY++ = p.m_Position.m_Y;
					*photonsZ++ = p.m_Position.m_Z;
				}

				// init memory structures and stuff ..
				StartAllocationStats();
				CUT_SAFE_CALL(cutStartTimer(m_ConstructionTimer));
				PreparekNNSearch(photons, photonsWidth, m_kNNBFQueryWidth, 3, static_cast<int>(m_MaxPhotons));
				CUT_SAFE_CALL(cutStopTimer(m_ConstructionTimer));
				EndAllocationStats();
				delete[] photons;
			}
			break;
	}

	float timeElapsed = cutGetTimerValue(m_ConstructionTimer);
	_tprintf(_T("Construction time: %f ms (%f s)\n"), timeElapsed, timeElapsed / 1000.0f);
}

Color PhotonMap::RadianceEstimate(const Math::Vector3& p_Point)
{
	using namespace Math;

	Color output(0.0f, 0.0f, 0.0f);

	LocateContext ctx(p_Point, m_MaxPhotons, m_QueryRadius * m_QueryRadius);
	LocatePhotons(ctx, 1);

	// sum irradiance from all the photons
	for(uintx i = 1; i <= ctx.m_PhotonsFound; ++i)
	{
		const Photon* p = ctx.m_Photons[i];
		Color power;
		power.FromRGBE(p->m_Power);

		//float weight = ConicFilter(Sqrt(ctx.m_SqDistances[i]), Sqrt(ctx.m_SqDistances[0]));
		float weight = GaussianFilter(ctx.m_SqDistances[i], ctx.m_SqDistances[0]);
		output += power * weight;
	}

	// estimate of density
	const float coneFilter = 1.1f;
	//const float tmp = 1.0f / ((1 - 2 / (3 * coneFilter)) * (PI * ctx.m_SqDistances[0]));
	const float tmp = 1.0f / (PI * ctx.m_SqDistances[0]);	
	output = output * tmp;
	return output;
}

bool PhotonMap::Store(const tchar* p_FileName, PhotonMapPtr p_Map)
{
	using namespace std;

	// open file
	ofstream stream(p_FileName, ios_base::binary | ios_base::out);
	if(!stream.is_open())
		return false;

	// write number of photons
	int32 numPhotons = static_cast<int32>(p_Map->m_PhotonMap.size());
	stream.write(reinterpret_cast<const char*>(&numPhotons), sizeof(numPhotons));

	// write bounding boxes
	stream.write(reinterpret_cast<const char*>(&p_Map->m_BBMin), sizeof(p_Map->m_BBMin));
	stream.write(reinterpret_cast<const char*>(&p_Map->m_BBMax), sizeof(p_Map->m_BBMax));

	// write photon map
	stream.write(reinterpret_cast<const char*>(&p_Map->m_PhotonMap[0]), sizeof(Photon) * numPhotons);
	stream.close();

	return true;
}

PhotonMapPtr PhotonMap::Load(const tchar* p_FileName, float p_MaxDist, uintx p_MaxPhotons, uintx p_Technique)
{
	using namespace std;

	ifstream stream(p_FileName, ios_base::binary | ios_base::in);
	if(!stream.is_open())
		return PhotonMapPtr();

	// read number of photons
	PhotonMapPtr pm = PhotonMapPtr(new PhotonMap(p_MaxDist, p_MaxPhotons, p_Technique));
	int32 numPhotons;
	stream.read(reinterpret_cast<char*>(&numPhotons), sizeof(numPhotons));

	// read bounding boxes
	stream.read(reinterpret_cast<char*>(&pm->m_BBMin), sizeof(pm->m_BBMin));
	stream.read(reinterpret_cast<char*>(&pm->m_BBMax), sizeof(pm->m_BBMax));

	// read photon map
	pm->m_PhotonMap.resize(numPhotons);
	stream.read(reinterpret_cast<char*>(&pm->m_PhotonMap[0]), sizeof(Photon) * numPhotons);
	stream.close();
	pm->m_NumPhotons = pm->m_PhotonMap.size() - 1;
	pm->m_NumPhotonsHalf = pm->m_NumPhotons / 2 - 1;

	_tprintf(_T("Loaded photon map with %d photons\n"), pm->m_NumPhotons);

	return pm;
}

void PhotonMap::BuildKdTree()
{
	using namespace std;

	// not necessarily good sc++l practice to do it this way but the algorithm
	// is based on using some pointer arithmetic so we use vector in a bad way
	// here for simplicity
	// the code is taken from the book Realistic Image Synthesis using Photon Mapping
	// by H. W. Jensen
	vector<Photon*> pa1(m_NumPhotons + 1);
	vector<Photon*> pa2(m_NumPhotons + 1);
	for(size_t i = 1; i <= m_NumPhotons; ++i)
		pa2[i] = &m_PhotonMap[i];

	BalanceSegment(pa1, pa2, 1, 1, static_cast<int>(m_NumPhotons));

	// reorganize balanced kd-tree (make a heap)
	intx d, j = 1, foo = 1;
	Photon foo_photon = m_PhotonMap[j];
	for(uintx i = 1; i <= m_NumPhotons; ++i)
	{
		d = pa1[j] - &(m_PhotonMap[0]);
		pa1[j] = NULL;
		if(d != foo)
			m_PhotonMap[j] = m_PhotonMap[d];
		else 
		{
			m_PhotonMap[j] = foo_photon;
			if(i < m_NumPhotons) 
			{
				for(; static_cast<uintx>(foo) <= m_NumPhotons; ++foo)
					if(pa1[foo] != NULL)
						break;
				foo_photon = m_PhotonMap[foo];
				j = foo;
			}
			continue;
		}
		j = d;
	}
	m_NumPhotonsHalf = m_NumPhotons / 2 - 1;
}

void PhotonMap::BalanceSegment(std::vector<Photon*>& p_Bal, std::vector<Photon*>& p_Org, const int p_Index, const int p_Start, const int p_End)
{
	// compute new median
	int median = 1;
	while((4 * median) <= (p_End - p_Start + 1))
		median += median;

	if((3 * median) <= (p_End - p_Start + 1)) 
	{
		median += median;
		median += p_Start - 1;
	}
	else
		median = p_End - median + 1;	

	// find axis to split along
	int axis = 2;
	if(((m_BBMax.m_X - m_BBMin.m_X) >  (m_BBMax.m_Y - m_BBMin.m_Y)) &&
		((m_BBMax.m_X - m_BBMin.m_X) >  (m_BBMax.m_Z - m_BBMin.m_Z)))
		axis = 0;
	else if((m_BBMax.m_Y - m_BBMin.m_Y) > (m_BBMax.m_Z - m_BBMin.m_Z))
		axis = 1;

	// partition photon block around the median
	MedianSplit(p_Org, p_Start, p_End, median, axis);

	p_Bal[p_Index] = p_Org[median];
	p_Bal[p_Index]->m_Flags = axis;

	// recursively balance the left and right block
	if(median > p_Start) 
	{
		// balance left segment
		if(p_Start < median - 1)
		{
			const float tmp = m_BBMax.m_Xyz[axis];
			m_BBMax.m_Xyz[axis] = p_Bal[p_Index]->m_Position.m_Xyz[axis];      
			BalanceSegment(p_Bal, p_Org, 2 * p_Index, p_Start, median - 1);
			m_BBMax.m_Xyz[axis] = tmp;
		}    
		else
			p_Bal[2 * p_Index] = p_Org[p_Start];
	}

	if(median < p_End)
	{
		// balance right segment
		if(median + 1 < p_End)
		{
			const float tmp = m_BBMin.m_Xyz[axis];
			m_BBMin.m_Xyz[axis] = p_Bal[p_Index]->m_Position.m_Xyz[axis];      
			BalanceSegment(p_Bal, p_Org, 2 * p_Index + 1, median + 1, p_End);
			m_BBMin.m_Xyz[axis] = tmp;
		}

		else
			p_Bal[2 * p_Index + 1] = p_Org[p_End];
	}
}

void PhotonMap::MedianSplit(std::vector<Photon*>& p_Photons, const int p_Start, const int p_End, const int p_Median, const int p_Axis)
{
	using namespace std;

	int left = p_Start;
	int right = p_End;

	while(right > left)
	{

		const float v = p_Photons[right]->m_Position.m_Xyz[p_Axis];
		int i = left - 1;
		int j = right;
		for (;;)
		{
			while(p_Photons[++i]->m_Position.m_Xyz[p_Axis] < v)
				;
			while(p_Photons[--j]->m_Position.m_Xyz[p_Axis] > v && j > left)
				;
			if(i >= j)
				break;
			swap(p_Photons[i], p_Photons[j]);
		}

		swap(p_Photons[i], p_Photons[right]);
		if(i >= p_Median)
			right = i - 1;
		if(i <= p_Median)
			left = i + 1;
	}
}

const Math::Vector3 PhotonMap::GetPhotonDir(const byte p_Theta, const byte p_Phi)
{
	return Math::Vector3(m_SinTheta[p_Theta] * m_CosPhi[p_Phi], m_SinTheta[p_Theta] * m_SinPhi[p_Phi], m_CosTheta[p_Theta]);
}

void PhotonMap::LocatePhotons(LocateContext& p_Context, const uintx p_Index)
{
	const Photon *p = &m_PhotonMap[p_Index];
	float dist1, dist2;

	if(p_Index < m_NumPhotonsHalf)
	{	
		dist1 = p_Context.m_Pos.m_Xyz[p->m_Flags] - p->m_Position.m_Xyz[p->m_Flags];
		
		// if dist1 is positive search right half-plane first
		if(dist1 > 0.0f) 
		{
			LocatePhotons(p_Context, 2 * p_Index + 1);
			if((dist1 * dist1) < p_Context.m_SqDistances[0])
				LocatePhotons(p_Context, 2 * p_Index);
		}
		else // dist1 is negative, search left half-plane first
		{
			LocatePhotons(p_Context, 2 * p_Index);
			if((dist1 * dist1) < p_Context.m_SqDistances[0])
				LocatePhotons(p_Context, 2 * p_Index + 1);
		}
	}

	// compute squared distance between current photon and np->pos
	dist2 = (p->m_Position - p_Context.m_Pos).LengthSqr();
	if(dist2 < p_Context.m_SqDistances[0])
	{
		// we found a photon  [:)] Insert it in the candidate list
		if(p_Context.m_PhotonsFound < p_Context.m_MaxPhotons) 
		{
			// heap is not full; use array
			p_Context.m_PhotonsFound++;
			p_Context.m_SqDistances[p_Context.m_PhotonsFound] = dist2;
			p_Context.m_Photons[p_Context.m_PhotonsFound] = p;
		}
		else 
		{
			// Do we need to build the heap?
			uintx j,parent;
			if(p_Context.m_BuildHeap) 
			{ 
				// Build heap
				float dst2;
				const Photon *phot;
				uintx half_found = (p_Context.m_PhotonsFound >> 1);
				for(intx k = half_found; k >= 1; k--)
				{
					parent = k;
					phot = p_Context.m_Photons[k];
					dst2 = p_Context.m_SqDistances[k];
					while(parent <= half_found)
					{
						j = parent + parent;
						if(j < p_Context.m_PhotonsFound && p_Context.m_SqDistances[j] < p_Context.m_SqDistances[j+1])
							j++;
						if(dst2 >= p_Context.m_SqDistances[j])
							break;
						p_Context.m_SqDistances[parent] = p_Context.m_SqDistances[j];
						p_Context.m_Photons[parent] = p_Context.m_Photons[j];
						parent = j;
					}
					p_Context.m_SqDistances[parent] = dst2;
					p_Context.m_Photons[parent] = phot;
				}
				p_Context.m_BuildHeap = 0;
			}

			// insert new photon into max heap
			// delete largest element, insert new and reorder the heap
			parent=1;
			j = 2;
			while(j <= p_Context.m_PhotonsFound)
			{
				if(j < p_Context.m_PhotonsFound && p_Context.m_SqDistances[j] < p_Context.m_SqDistances[j+1])
					j++;
				if(dist2 > p_Context.m_SqDistances[j])
					break;
				p_Context.m_SqDistances[parent] = p_Context.m_SqDistances[j];
				p_Context.m_Photons[parent] = p_Context.m_Photons[j];
				parent = j;
				j += j;
			}
			if (dist2 < p_Context.m_SqDistances[parent])
			{
				p_Context.m_SqDistances[parent] = dist2;
				p_Context.m_Photons[parent] = p;
			}

			p_Context.m_SqDistances[0] = p_Context.m_SqDistances[1];
		}
	}
}

void PhotonMap::DeferredEstimate(RayTraceRecorderPtr p_Recorder)
{
	using namespace Math;

	// select the right technique
	switch(m_Technique)
	{
		case SWKdTree:
			{
				CUT_SAFE_CALL(cutStartTimer(m_EstimateTimer));
				CUT_SAFE_CALL(cutStartTimer(m_PhotonGatherTimer));
				for(size_t i = 0; i < p_Recorder->m_BufferOut.size(); ++i)
					(*p_Recorder->m_BufferOut[i]) += (p_Recorder->m_Factors[i] * 
						RadianceEstimate(p_Recorder->m_QueryPoints[i]));
				CUT_SAFE_CALL(cutStopTimer(m_PhotonGatherTimer));
				CUT_SAFE_CALL(cutStopTimer(m_EstimateTimer));
			}
			break;

		case GPUKdTree:
			DeferredEstimateKdTree(p_Recorder);
			break;
			
		case GPUHash:
			DeferredEstimateHash(p_Recorder);
			break;
				
		case GPUBF:
		default:
			DeferredEstimateBF(p_Recorder);
			break;
	}

	float timeElapsed = cutGetTimerValue(m_EstimateTimer);
	_tprintf(_T("Total estimation done in %f ms (%f s)\n"), timeElapsed, timeElapsed / 1000.0f);
	timeElapsed = cutGetTimerValue(m_PhotonGatherTimer);
	_tprintf(_T("Photon gathering done in %f ms (%f s)\n"), timeElapsed, timeElapsed / 1000.0f);
}

void PhotonMap::DeferredEstimateBF(RayTraceRecorderPtr p_Recorder)
{
	using namespace Math;

	// allocate memory for query and results
	int queryIterations = static_cast<int>(p_Recorder->m_BufferOut.size() / m_kNNBFQueryWidth);
	if(p_Recorder->m_BufferOut.size() % m_kNNBFQueryWidth) queryIterations++;
	float* queries = new float[p_Recorder->m_BufferOut.size() * 3];
	float** sqDist = new float*[queryIterations];
	int** indexes = new int*[queryIterations];

	// allocate space for results
	int photonsWidth = static_cast<int>(m_PhotonMap.size() - 1);
	for(int i = 0; i < queryIterations; ++i)
	{
		sqDist[i] = new float[m_kNNBFQueryWidth * m_MaxPhotons];
		indexes[i] = new int[m_kNNBFQueryWidth * m_MaxPhotons];
	}

	CUT_SAFE_CALL(cutStartTimer(m_EstimateTimer));

	// get queries into right format
	for(int j = 0; j < queryIterations; ++j)
	{
		int baseIdx = j * m_kNNBFQueryWidth;
		int actual_query_widht = static_cast<int>(__min(m_kNNBFQueryWidth, p_Recorder->m_BufferOut.size() - baseIdx));

		float* queryX = queries + baseIdx * 3;
		float* queryY = queryX + actual_query_widht;
		float* queryZ = queryY + actual_query_widht;

		for(int i = 0; i < actual_query_widht; ++i)
		{
			const Math::Vector3& v = p_Recorder->m_QueryPoints[baseIdx + i];
			*queryX++ = v.m_X;
			*queryY++ = v.m_Y;
			*queryZ++ = v.m_Z;
		}
	}

	// perform knn search
	unsigned int curRecPos = 0;
	CUT_SAFE_CALL(cutStartTimer(m_PhotonGatherTimer));
	for(int i = 0, j = 0; i < p_Recorder->m_QueryPoints.size(); i += m_kNNBFQueryWidth, ++j)
	{
		int actual_query_width = static_cast<int>(__min(m_kNNBFQueryWidth, p_Recorder->m_BufferOut.size() - i));
		knn(photonsWidth, &queries[i * 3], actual_query_width, 3, static_cast<int>(m_MaxPhotons), sqDist[j], indexes[j]);
	}
	CUT_SAFE_CALL(cutStopTimer(m_PhotonGatherTimer));

	// process results
	unsigned int queryBase = 0;
	unsigned int curPixel = 0;
	for(int i = 0; i < queryIterations; ++i)
	{
		int actual_query_width = static_cast<int>(__min(m_kNNBFQueryWidth, p_Recorder->m_BufferOut.size() - queryBase));
		for(int j = 0; j < actual_query_width; ++j)
		{
			// calc some parameters
			int* curIndex = indexes[i]+ j;
			float* curSqDist = sqDist[i] + j;
			float sqMaxDist = sqDist[i][j + actual_query_width * 19];

			// sum irradiance from all the photons
			Color output(0.0f, 0.0f, 0.0f);				
			for(uintx k = 0; k < m_MaxPhotons; ++k)
			{
				const Photon& p = m_PhotonMap[*curIndex];
				Color power;
				power.FromRGBE(p.m_Power);

				//float weight = ConicFilter(Sqrt(ctx.m_SqDistances[i]), Sqrt(ctx.m_SqDistances[0]));
				float weight = GaussianFilter(*curSqDist, sqMaxDist);
				output += power * weight;

				curIndex += actual_query_width;
				curSqDist += actual_query_width;
			}

			// estimate of density
			const float tmp = 1.0f / (PI * sqMaxDist);
			output = output * tmp;
			(*p_Recorder->m_BufferOut[curPixel]) += p_Recorder->m_Factors[curPixel] * output;
			curPixel++;
		}
		queryBase += actual_query_width;
	}

	CUT_SAFE_CALL(cutStopTimer(m_EstimateTimer));

	CleanUp();
	delete[] queries;
	for(int i = 0; i < queryIterations; ++i)
	{
		delete[] sqDist[i];
		delete[] indexes[i];
	}
	delete[] sqDist;
	delete[] indexes;
}

void PhotonMap::DeferredEstimateKdTree(RayTraceRecorderPtr p_Recorder)
{
	using namespace Math;

	// alloc memory for query and results
	float* queries = new float[p_Recorder->m_BufferOut.size() * 4];
	float* sqDistances = new float[sizeof(float) * m_MaxPhotons * p_Recorder->m_BufferOut.size()];
	byte* powers = new byte[sizeof(byte) * 4 * m_MaxPhotons * p_Recorder->m_BufferOut.size()];

	CUT_SAFE_CALL(cutStartTimer(m_EstimateTimer));

	// get queries into right format
	float* curQueryPt = queries;
	for(int i = 0; i < p_Recorder->m_BufferOut.size(); ++i)
	{
		const Math::Vector3& v = p_Recorder->m_QueryPoints[i];
		*curQueryPt++ = v.m_X;
		*curQueryPt++ = v.m_Y;
		*curQueryPt++ = v.m_Z;
		curQueryPt++;	// last element is unused
	}

	// search nearest neighbors
	CUT_SAFE_CALL(cutStartTimer(m_PhotonGatherTimer));
	SearchkNN(m_KdTree, static_cast<unsigned int>(m_PhotonMap.size() - 1), queries, 
		static_cast<unsigned int>(p_Recorder->m_BufferOut.size()), 
		static_cast<unsigned int>(m_MaxPhotons), m_QueryRadius, sqDistances, powers);
	CUT_SAFE_CALL(cutStopTimer(m_PhotonGatherTimer));

	// process results
	for(int i = 0; i < p_Recorder->m_BufferOut.size(); ++i)
	{
		float* curSqDist = sqDistances + i * m_MaxPhotons;
		byte* curPowers = powers + i * m_MaxPhotons * 4;

		// sum irradiance from all the photons
		Color output(0.0f, 0.0f, 0.0f);
		float maxSqDist = __min(m_QueryRadius, curSqDist[m_MaxPhotons - 1]);
		for(uintx k = 0; k < m_MaxPhotons && curSqDist[k] < maxSqDist; ++k)
		{
			Color power;
			power.FromRGBE(&curPowers[k * 4]);

			//float weight = ConicFilter(Sqrt(ctx.m_SqDistances[i]), Sqrt(ctx.m_SqDistances[0]));
			float weight = GaussianFilter(curSqDist[k], maxSqDist);
			output += power * weight;
		}

		// estimate of density
		const float tmp = 1.0f / (PI * maxSqDist);
		output = output * tmp;
		(*p_Recorder->m_BufferOut[i]) += p_Recorder->m_Factors[i] * output;
	}

	CUT_SAFE_CALL(cutStopTimer(m_EstimateTimer));

	// cleanup
	CleanUpKdTree(m_KdTree);
	delete[] queries;
	delete[] sqDistances;
	delete[] powers;
}

void PhotonMap::DeferredEstimateHash(RayTraceRecorderPtr p_Recorder)
{
	using namespace Math;

	// alloc memory for query and the results
	float* sqDistances = new float[sizeof(float) * (m_MaxPhotons) * p_Recorder->m_BufferOut.size()];
	byte* powers = new byte[sizeof(byte) * 4 * (m_MaxPhotons) * p_Recorder->m_BufferOut.size()];
	float* queries = new float[p_Recorder->m_BufferOut.size() * 4];

	CUT_SAFE_CALL(cutStartTimer(m_EstimateTimer));

	// get queries into right format
	float* curQueryPt = queries;
	for(int i = 0; i < p_Recorder->m_BufferOut.size(); ++i)
	{
		const Math::Vector3& v = p_Recorder->m_QueryPoints[i];
		*curQueryPt++ = v.m_X;
		*curQueryPt++ = v.m_Y;
		*curQueryPt++ = v.m_Z;
		curQueryPt++;	// last element is unused
	}

	// search nearest neighbors
	CUT_SAFE_CALL(cutStartTimer(m_PhotonGatherTimer));
	HashSearch(m_HashMap, queries, static_cast<unsigned int>(p_Recorder->m_BufferOut.size()), 
		m_QueryRadius, static_cast<unsigned int>(m_MaxPhotons), sqDistances, powers);
	CUT_SAFE_CALL(cutStopTimer(m_PhotonGatherTimer));

	// process results
	for(int i = 0; i < p_Recorder->m_BufferOut.size(); ++i)
	{
		float* curSqDist = sqDistances + i * m_MaxPhotons;
		byte* curPowers = powers + i * m_MaxPhotons * 4;
		float maxSqDist = __min(m_QueryRadius, curSqDist[m_MaxPhotons - 1]);

		// sum irradiance from all the photons
		Color output(0.0f, 0.0f, 0.0f);
		for(uintx k = 0; k < m_MaxPhotons && curSqDist[k] < maxSqDist; ++k)
		{
			Color power;
			power.FromRGBE(&curPowers[k * 4]);

			//float weight = ConicFilter(Sqrt(ctx.m_SqDistances[i]), Sqrt(ctx.m_SqDistances[0]));
			float weight = GaussianFilter(curSqDist[k], maxSqDist);
			output += power * weight;
		}

		// estimate of density
		const float tmp = 1.0f / (PI * maxSqDist);
		output = output * tmp;
		(*p_Recorder->m_BufferOut[i]) += p_Recorder->m_Factors[i] * output;
	}

	CUT_SAFE_CALL(cutStopTimer(m_EstimateTimer));

	delete[] queries;
	delete[] sqDistances;
	delete[] powers;
	CleanUpHash(m_HashMap);
}

}

